import sys
import os
from PIL import Image, ImageQt
import numpy as np
import cv2
from PyQt5.QtWidgets import (QApplication, QMainWindow, QPushButton, QLabel,
                             QVBoxLayout, QHBoxLayout, QFileDialog, QListWidget,
                             QWidget, QMessageBox, QGroupBox, QFormLayout, QSpinBox,
                             QDoubleSpinBox, QProgressBar)
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from ultralytics import YOLO


class DetectionThread(QThread):
    """检测线程，避免UI卡顿"""
    finished = pyqtSignal(object, int, dict)
    progress = pyqtSignal(int)

    def __init__(self, model_path, image_path, conf_threshold, iou_threshold):
        super().__init__()
        self.model_path = model_path
        self.image_path = image_path
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold

    def run(self):
        try:
            # 加载模型
            model = YOLO(self.model_path)

            # 执行检测
            results = model.predict(
                source=self.image_path,
                conf=self.conf_threshold,
                iou=self.iou_threshold,
                verbose=False
            )

            # 获取检测结果
            result = results[0]
            detected_image = result.plot()

            # 计算羊的数量
            sheep_count = len(result.boxes)

            # 收集模型指标
            metrics = {
                "confidence_min": float(result.boxes.conf.min()) if sheep_count > 0 else 0,
                "confidence_max": float(result.boxes.conf.max()) if sheep_count > 0 else 0,
                "confidence_avg": float(result.boxes.conf.mean()) if sheep_count > 0 else 0
            }

            self.finished.emit(detected_image, sheep_count, metrics)
        except Exception as e:
            self.finished.emit(None, 0, {"error": str(e)})


class SheepDetectorGUI(QMainWindow):
    """湖羊数量检测主界面"""

    def __init__(self):
        super().__init__()
        self.initUI()

        # 初始化变量
        self.images = []  # 存储导入的图片路径
        self.current_image_index = -1  # 当前选中的图片索引
        self.model_path = ""  # 模型路径

    def initUI(self):
        """初始化界面"""
        self.setWindowTitle('湖羊数量检测系统')
        self.setGeometry(100, 100, 1200, 800)

        # 创建主布局
        main_layout = QHBoxLayout()

        # 左侧控制面板
        control_panel = QVBoxLayout()

        # 模型加载部分
        model_group = QGroupBox("模型设置")
        model_layout = QVBoxLayout()

        self.model_path_label = QLabel("未选择模型")
        model_layout.addWidget(self.model_path_label)

        model_btn_layout = QHBoxLayout()
        self.load_model_btn = QPushButton("加载模型")
        self.load_model_btn.clicked.connect(self.load_model)
        model_btn_layout.addWidget(self.load_model_btn)

        model_layout.addLayout(model_btn_layout)
        model_group.setLayout(model_layout)
        control_panel.addWidget(model_group)

        # 图片导入部分
        image_group = QGroupBox("图片管理")
        image_layout = QVBoxLayout()

        self.image_list = QListWidget()
        self.image_list.itemClicked.connect(self.on_image_selected)
        image_layout.addWidget(self.image_list)

        btn_layout = QHBoxLayout()
        self.import_btn = QPushButton("导入图片")
        self.import_btn.clicked.connect(self.import_images)
        self.clear_btn = QPushButton("清空列表")
        self.clear_btn.clicked.connect(self.clear_images)

        btn_layout.addWidget(self.import_btn)
        btn_layout.addWidget(self.clear_btn)
        image_layout.addLayout(btn_layout)

        image_group.setLayout(image_layout)
        control_panel.addWidget(image_group)

        # 参数设置部分
        params_group = QGroupBox("检测参数")
        params_layout = QFormLayout()

        self.conf_threshold = QDoubleSpinBox()
        self.conf_threshold.setRange(0.1, 0.9)
        self.conf_threshold.setValue(0.5)
        self.conf_threshold.setSingleStep(0.05)

        self.iou_threshold = QDoubleSpinBox()
        self.iou_threshold.setRange(0.1, 0.9)
        self.iou_threshold.setValue(0.7)
        self.iou_threshold.setSingleStep(0.05)

        params_layout.addRow("置信度阈值:", self.conf_threshold)
        params_layout.addRow("IOU阈值:", self.iou_threshold)

        params_group.setLayout(params_layout)
        control_panel.addWidget(params_group)

        # 检测按钮
        self.detect_btn = QPushButton("开始检测")
        self.detect_btn.clicked.connect(self.detect)
        self.detect_btn.setEnabled(False)
        control_panel.addWidget(self.detect_btn)

        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setTextVisible(True)
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(0)
        control_panel.addWidget(self.progress_bar)

        # 将控制面板添加到主布局
        control_panel_widget = QWidget()
        control_panel_widget.setLayout(control_panel)
        control_panel_widget.setMaximumWidth(300)
        main_layout.addWidget(control_panel_widget)

        # 右侧显示面板
        display_panel = QVBoxLayout()

        # 图片显示区域
        self.image_display = QLabel("请导入图片并选择一张进行检测")
        self.image_display.setAlignment(Qt.AlignCenter)
        self.image_display.setMinimumSize(800, 600)
        self.image_display.setStyleSheet("border: 1px solid #ccc;")
        display_panel.addWidget(self.image_display)

        # 结果显示区域
        results_group = QGroupBox("检测结果")
        results_layout = QFormLayout()

        self.sheep_count_label = QLabel("羊数量: --")
        self.conf_min_label = QLabel("最小置信度: --")
        self.conf_max_label = QLabel("最大置信度: --")
        self.conf_avg_label = QLabel("平均置信度: --")

        results_layout.addRow(self.sheep_count_label)
        results_layout.addRow(self.conf_min_label)
        results_layout.addRow(self.conf_max_label)
        results_layout.addRow(self.conf_avg_label)

        results_group.setLayout(results_layout)
        display_panel.addWidget(results_group)

        # 将显示面板添加到主布局
        main_layout.addLayout(display_panel)

        # 设置主窗口中心部件
        central_widget = QWidget()
        central_widget.setLayout(main_layout)
        self.setCentralWidget(central_widget)

    def load_model(self):
        """加载YOLOv8模型"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择模型文件", "", "模型文件 (*.pt)"
        )

        if file_path:
            self.model_path = file_path
            self.model_path_label.setText(os.path.basename(file_path))
            QMessageBox.information(self, "成功", "模型加载成功")

            # 如果有图片，启用检测按钮
            if self.images:
                self.detect_btn.setEnabled(True)

    def import_images(self):
        """导入图片"""
        file_paths, _ = QFileDialog.getOpenFileNames(
            self, "选择图片", "", "图片文件 (*.png *.jpg *.jpeg *.bmp)"
        )

        if file_paths:
            self.images.extend(file_paths)
            self.update_image_list()

            # 如果已加载模型，启用检测按钮
            if self.model_path:
                self.detect_btn.setEnabled(True)

    def clear_images(self):
        """清空图片列表"""
        self.images = []
        self.image_list.clear()
        self.current_image_index = -1
        self.image_display.setText("请导入图片并选择一张进行检测")
        self.reset_results()

        # 禁用检测按钮
        self.detect_btn.setEnabled(False)

    def update_image_list(self):
        """更新图片列表显示"""
        self.image_list.clear()
        for img_path in self.images:
            self.image_list.addItem(os.path.basename(img_path))

    def on_image_selected(self, item):
        """图片选择事件处理"""
        index = self.image_list.row(item)
        self.current_image_index = index
        self.display_image(self.images[index])
        self.reset_results()

    def display_image(self, image_path):
        """显示图片"""
        pixmap = QPixmap(image_path)
        scaled_pixmap = pixmap.scaled(
            self.image_display.width(),
            self.image_display.height(),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        )
        self.image_display.setPixmap(scaled_pixmap)

    def reset_results(self):
        """重置结果显示"""
        self.sheep_count_label.setText("羊数量: --")
        self.conf_min_label.setText("最小置信度: --")
        self.conf_max_label.setText("最大置信度: --")
        self.conf_avg_label.setText("平均置信度: --")

    def detect(self):
        """执行检测"""
        if self.current_image_index == -1:
            QMessageBox.warning(self, "警告", "请先选择一张图片")
            return

        if not self.model_path:
            QMessageBox.warning(self, "警告", "请先加载模型")
            return

        # 获取当前图片路径
        image_path = self.images[self.current_image_index]

        # 创建并启动检测线程
        self.detection_thread = DetectionThread(
            self.model_path,
            image_path,
            self.conf_threshold.value(),
            self.iou_threshold.value()
        )
        self.detection_thread.finished.connect(self.on_detection_finished)
        self.detection_thread.progress.connect(self.update_progress)
        self.detection_thread.start()

        # 更新UI状态
        self.detect_btn.setEnabled(False)
        self.progress_bar.setValue(0)

    def update_progress(self, value):
        """更新进度条"""
        self.progress_bar.setValue(value)

    def on_detection_finished(self, result_image, sheep_count, metrics):
        """检测完成后的回调"""
        # 更新UI状态
        self.detect_btn.setEnabled(True)
        self.progress_bar.setValue(100)

        # 处理检测结果
        if result_image is None:
            if "error" in metrics:
                QMessageBox.critical(self, "错误", f"检测失败: {metrics['error']}")
            else:
                QMessageBox.critical(self, "错误", "检测失败，未知错误")
            return

        # 转换结果图像为Qt格式并显示
        if isinstance(result_image, np.ndarray):
            result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)
            h, w, ch = result_image.shape
            bytes_per_line = ch * w
            q_img = QImage(result_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
            pixmap = QPixmap.fromImage(q_img)
        else:
            pixmap = QPixmap.fromImage(ImageQt.ImageQt(result_image))

        scaled_pixmap = pixmap.scaled(
            self.image_display.width(),
            self.image_display.height(),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        )
        self.image_display.setPixmap(scaled_pixmap)

        # 更新结果显示
        self.sheep_count_label.setText(f"羊数量: {sheep_count}")
        self.conf_min_label.setText(f"最小置信度: {metrics['confidence_min']:.4f}")
        self.conf_max_label.setText(f"最大置信度: {metrics['confidence_max']:.4f}")
        self.conf_avg_label.setText(f"平均置信度: {metrics['confidence_avg']:.4f}")


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = SheepDetectorGUI()
    window.show()
    sys.exit(app.exec_())